import os
from typing import Annotated

from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
from langchain_core.runnables.graph import MermaidDrawMethod
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI


from config.model_config import get_chat_openai_zhipu, get_chat_openai_zhipu_flash_250414


class State(TypedDict):
    messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)

# 实际请以官方文档为准
# llm = get_chat_openai_zhipu()
llm = get_chat_openai_zhipu_flash_250414()

def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}

graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()

##
# 支持的流模式
# LangGraph
# 支持五种核心流模式，每种模式解决不同的场景需求：
#
# 模式
# 描述
# values 流式传输每一步后的完整状态值
# updates 流式传输每一步的状态更新内容，同一步的多次更新会分开流式传输
# custom 从节点内部流式传输自定义数据的
# messages  流式传输
# LLM  调用的令牌及元数据，适用于实时展示模型生成过程
# debug
# 流式传输执行过程中的详细调试信息，包含节点名称和完整状态

def stream_graph_updates(user_input: str):
    print("Assistant:", end=' ')
    for step_state  in graph.stream({"messages": [{"role": "user", "content": user_input}]}, stream_mode="updates"):
        print(step_state)
    print("\n")
while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "一句话介绍langchain和LangGraph"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break
